import torch 
from PIL import Image
import numpy as np


class High(torch.nn.Module):
    def __init__(self, M, N) -> None:
        super().__init__()
           # 提取高频成分
        self.mask_high_pass = torch.zeros((M, N))
        d = 30  # 高频阈值
        crow, ccol = M // 2, N // 2
        self.mask_high_pass[crow-d:crow+d, ccol-d:ccol+d] = 1


    def forward(self, img):
        freq_ = torch.fft.fft2(img)
        freq = torch.fft.fftshift(freq_)
        high_pass = self.mask_high_pass * freq  
        ishift = torch.fft.ifftshift(high_pass)   
        ori_img = torch.fft.ifft2(ishift)
        return torch.abs(ori_img)
    

img = Image.open('lena.png') 
img = np.array(img).astype(np.float32)
img = torch.tensor(img)
img = img.permute(2, 0, 1)
img.requires_grad = True
layer = High(512, 512)
c = layer(img)

c_np = c.detach().numpy()
c_np = c_np.transpose(1, 2, 0)
c_np = (c_np-c_np.min())/(c_np.max()-c_np.min())
c_np = c_np * 255
c_np = c_np.astype(np.uint8)
I = Image.fromarray(c_np)
I.show()


label = torch.ones_like(img)
loss = (label - img).sum()
loss.backward()
print('end')